#!/usr/bin/env python3
import time
import json
import matplotlib.pyplot as plt
import pathmagic
import random
import math
import sys
sys.path.append(r"C:\Users\Admin\Desktop\CBBA_Python_main\CBBA_Python_main")


with pathmagic.context():
    from lib.CBBA import CBBA
    from lib.WorldInfo import WorldInfo
    import lib.HelperLibrary as HelperLibrary
    from lib.Obstacle import Obstacle


# RRT相关函数
class Node:
    def __init__(self, x, y, z):
        self.x = x
        self.y = y
        self.z = z
        self.parent = None

def distance(node1, node2):
    return ((node1.x - node2.x)**2 + (node1.y - node2.y)**2 + (node1.z - node2.z)**2) ** 0.5

def sample_free_space(x_range, y_range, z_range):
    return Node(random.uniform(*x_range), random.uniform(*y_range), random.uniform(*z_range))

def nearest_node(nodes, sample):
    return min(nodes, key=lambda n: distance(n, sample))

def steer(from_node, to_node, extend_length=1.0):
    new_node = Node(from_node.x, from_node.y, from_node.z)
    if distance(from_node, to_node) < extend_length:
        new_node.x = to_node.x
        new_node.y = to_node.y
        new_node.z = to_node.z
    else:
        theta_x = math.atan2(to_node.x - from_node.x, distance(from_node, to_node))
        theta_y = math.atan2(to_node.y - from_node.y, distance(from_node, to_node))
        theta_z = math.atan2(to_node.z - from_node.z, distance(from_node, to_node))
        new_node.x += extend_length * math.cos(theta_x)
        new_node.y += extend_length * math.cos(theta_y)
        new_node.z += extend_length * math.cos(theta_z)
    new_node.parent = from_node
    return new_node

def collision_free(new_node, obstacle_list):
    for obs in obstacle_list:
        if (obs.x_min <= new_node.x <= obs.x_max and
            obs.y_min <= new_node.y <= obs.y_max and
            obs.z_min <= new_node.z <= obs.z_max):
            return False
    return True

def rrt(start, goal, x_range, y_range, z_range, obstacle_list, max_iter=1000, step_size=1.0):
    nodes = [start]
    for _ in range(max_iter):
        sample = sample_free_space(x_range, y_range, z_range)
        nearest = nearest_node(nodes, sample)
        new_node = steer(nearest, sample, step_size)
        if collision_free(new_node, obstacle_list):
            nodes.append(new_node)
            if distance(new_node, goal) < step_size:
                goal.parent = new_node
                nodes.append(goal)
                break
    return nodes

def reconstruct_path(goal):
    path = []
    node = goal
    while node is not None:
        path.append((node.x, node.y, node.z))
        node = node.parent
    return list(reversed(path))

if __name__ == "__main__":
    # a json configuration file
    config_file_name = "config_example_03.json"
    # Read the configuration from the json file
    json_file = open(config_file_name)
    config_data = json.load(json_file)

    # create a world, each list is [min, max] coordinates for x,y,z axis
    WorldInfoTest = WorldInfo([-2.0, 2.5], [-1.5, 5.5], [0.0, 20.0])

    # 创建障碍物
    obstacle1 = Obstacle([0.0, 0.0, 0.0], [1.0, 1.0, 1.0])  # 示例障碍物
    obstacle2 = Obstacle([1.5, 2.0, 0.0], [2.0, 2.5, 1.0])  # 另一个示例障碍物

    # 将障碍物添加到 WorldInfo 对象中
    WorldInfoTest.add_obstacle(obstacle1)
    WorldInfoTest.add_obstacle(obstacle2)

    # create a list of Agent(s) and Task(s)
    num_agents = 5
    num_tasks = 20
    max_depth = num_tasks
    AgentList, TaskList = HelperLibrary.create_agents_and_tasks_homogeneous(
        num_agents, num_tasks, WorldInfoTest, config_data)

    # create a CBBA solver
    CBBA_solver = CBBA(config_data)

    t_start = time.time()

    # solve, no time window
    path_list, _ = CBBA_solver.solve(AgentList, TaskList, WorldInfoTest, max_depth, time_window_flag=False)

    t_end = time.time()
    t_used = t_end - t_start
    print("Time used [sec]: ", t_used)

    # the output is CBBA_solver.path_list or path_list
    print("bundle_list")
    print(CBBA_solver.bundle_list)
    print("path_list")
    print(path_list)

    # 在每个代理的任务路径上应用RRT算法进行避障
    for agent_id, paths in enumerate(path_list):
        for i in range(len(paths) - 1):
            start = Node(TaskList[paths[i]].x,TaskList[paths[i]].y,TaskList[paths[i]].z)
            goal = Node(TaskList[paths[i + 1]].x,TaskList[paths[i + 1]].y,TaskList[paths[i + 1]].z)
            # start = Node(paths[i][0], paths[i][1], paths[i][2])
            # goal = Node(paths[i + 1][0], paths[i + 1][1], paths[i + 1][2])
            rrt_nodes = rrt(start, goal, WorldInfoTest.limit_x, WorldInfoTest.limit_y, WorldInfoTest.limit_z, WorldInfoTest.obstacles)
            if rrt_nodes:
                path_list[agent_id][i:i+2] = reconstruct_path(rrt_nodes[-1])

    # plot without time window
    CBBA_solver.plot_assignment_without_timewindow()
    plt.show()